#!/usr/bin/env python

from matplotlib import pyplot
import numpy

pathd = {'no_cut': 'no_cut_1p1s_yields.dat',
         'atl1abv15': 'nabv15_1p1s_yields.dat',
         'atl2abv15': 'atl2abv15_1p1s_yields.dat',
         }

datad = {}
xd    = {}
dxd   = {}
yd    = {}
dyd   = {}

for key in pathd:
    datad[key] = numpy.loadtxt(pathd[key], skiprows=5, comments='END').T

    xd[key] = datad[key][0]
    yd[key] = datad[key][1]
    dxd[key] = datad[key][2]
    dyd[key] = datad[key][3]

keys = ['no_cut', 'atl1abv15', 'atl2abv15']
titles = dict(zip(keys, [r'No track cuts', 
                    r'At least 1 track w/ $|p| > 1.5 GeV$',
                    r'At least 2 track w/ $|p| > 1.5 GeV$',
                    ]))

no_cut_ylim = None
for i, k in enumerate(keys):
    pyplot.errorbar(xd[k], yd[k], xerr=dxd[k], yerr=dyd[k], fmt='r.')

    pyplot.xlabel(r"$E_{\gamma}^*$ (GeV)")

    if k == 'no_cut':
        pyplot.ylabel(r"$h_b(1P)$ yield per 5 MeV")
        no_cut_ylim = pyplot.ylim()
    else:
        pyplot.ylim(no_cut_ylim)

    pyplot.title(titles[k])
    
    pyplot.savefig(k + '.png')
    pyplot.clf()


